import numpy as np
import torch

from ..base_bbox_coder import BaseBBoxCoder
from mmdet.core.bbox.transforms_obb import regular_theta, regular_obb
from mmdet.core.bbox.builder import BBOX_CODERS
import pdb
from math import pi


@BBOX_CODERS.register_module()
class AngleCoder(BaseBBoxCoder):

    def __init__(self,
                 theta_norm=True,
                 target_means=(0., 0., 0., 0., 0.),
                 target_stds=(1., 1., 1., 1., 1.)):
        super(AngleCoder, self).__init__()
        self.theta_norm = theta_norm
        self.means = target_means
        self.stds = target_stds

    def encode(self, bboxes, gt_bboxes):
        assert bboxes.size(0) == gt_bboxes.size(0)
        assert bboxes.size(-1) == 4
        assert gt_bboxes.size(-1) == 5
        encoded_bboxes = obb2delta(bboxes, gt_bboxes, self.theta_norm, self.means, self.stds)
        return encoded_bboxes

    def decode(self,
               bboxes,
               pred_bboxes,
               pred_angles,
               max_shape=None,
               wh_ratio_clip=16 / 1000):
        assert pred_bboxes.size(0) == bboxes.size(0)
        decoded_bboxes = delta2obb(bboxes, pred_bboxes, pred_angles, self.theta_norm,
                                   self.means, self.stds, wh_ratio_clip)

        return decoded_bboxes


def obb2delta(proposals, gt, theta_norm=True, means=(0., 0., 0., 0., 0.), stds=(1., 1., 1., 1., 1.)):
    proposals = proposals.float()  # proposal:(n,4) (x1,y1,x2,y2)
    gt = gt.float()  #
    px = (proposals[..., 0] + proposals[..., 2]) * 0.5  # px = (x1+x2)/2
    py = (proposals[..., 1] + proposals[..., 3]) * 0.5  # py = (y1+y2)/2
    pw = proposals[..., 2] - proposals[..., 0]  # pw = w
    ph = proposals[..., 3] - proposals[..., 1]  # ph = h
    gx, gy, gw, gh, gtheta = gt.unbind(dim=-1)  # gtheta属于 [-pi/2,pi/2)
    w_angle = torch.where(gw / gh >= 1.5, 1, 1)
    dtheta1 = regular_theta(gtheta)  # (gtheta+pi/2)%pi - pi/2          # (n)
    dtheta2 = regular_theta(gtheta + pi / 2)  # (gtheta+pi)% pi -pi/2    # (n)
    abs_dtheta1 = torch.abs(dtheta1)
    abs_dtheta2 = torch.abs(dtheta2)

    gw_regular = torch.where(abs_dtheta1 < abs_dtheta2, gw, gh)  # gw_regular (n) # (-pi/4,pi/4):gw 否则gh
    gh_regular = torch.where(abs_dtheta1 < abs_dtheta2, gh, gw)  # gh_regular (n) # (-pi/4,pi/4):gh 否则gw
    # dtheta = torch.where(abs_dtheta1 < abs_dtheta2, dtheta1,
    #                      dtheta2)  # dtheta (n)     # (-pi/4,pi/4):gtheta (pi/4,pi/2):pi/2-gtheta (-pi/2,-pi/4): pi/2+gtheta  # 这里才最终得到所谓的
    dx = (gx - px) / pw  # 找到物体的中心x点相对于pw的偏移
    dy = (gy - py) / ph  # 找到物体中的中心y点相对于py的偏移
    dw = torch.log(gw_regular / pw)  # w上的偏移相当于gt中的w,log(w/pw)
    dh = torch.log(gh_regular / ph)  # h上的偏移相当于gt中的h,log(h,/pw)

    # if theta_norm:
    #     dtheta /= 2 * pi
    deltas = torch.stack([dx, dy, dw, dh], dim=-1)

    means = deltas.new_tensor(means).unsqueeze(0)
    stds = deltas.new_tensor(stds).unsqueeze(0)
    deltas = deltas.sub_(means).div_(stds)
    deltas = torch.cat([deltas, gtheta.unsqueeze(-1), w_angle.unsqueeze(-1)], dim=-1)
    return deltas


def delta2obb(proposals,  # (n,4)
              deltas,  # (n,4)
              angle_preds,  # (n,60)
              theta_norm=True,
              means=(0., 0., 0., 0., 0.),
              stds=(1., 1., 1., 1., 1.),
              wh_ratio_clip=16 / 1000):
    means = deltas.new_tensor(means).repeat(1, deltas.size(1) // 4)
    stds = deltas.new_tensor(stds).repeat(1, deltas.size(1) // 4)
    # print(deltas.shape, stds.shape, means.shape)
    denorm_deltas = deltas * stds + means

    dx = denorm_deltas[:, 0::4]  # (n)
    dy = denorm_deltas[:, 1::4]
    dw = denorm_deltas[:, 2::4]
    dh = denorm_deltas[:, 3::4]
    # dtheta = denorm_deltas[:, 4::5]
    # if theta_norm:
    #     dtheta *= 2 * pi
    max_ratio = np.abs(np.log(wh_ratio_clip))
    dw = dw.clamp(min=-max_ratio, max=max_ratio)
    dh = dh.clamp(min=-max_ratio, max=max_ratio)

    px = ((proposals[:, 0] + proposals[:, 2]) * 0.5).unsqueeze(1).expand_as(dx)  # (n,
    py = ((proposals[:, 1] + proposals[:, 3]) * 0.5).unsqueeze(1).expand_as(dy)  # (n,
    pw = (proposals[:, 2] - proposals[:, 0]).unsqueeze(1).expand_as(dw)  #
    ph = (proposals[:, 3] - proposals[:, 1]).unsqueeze(1).expand_as(dh)

    gx = px + pw * dx
    gy = py + ph * dy
    gw = pw * dw.exp()
    gh = ph * dh.exp()
    # gtheta = dtheta

    gtheta = torch.argmax(angle_preds, dim=-1) * pi / 60 - pi / 2
    dtheta1 = regular_theta(gtheta)  # (gtheta+pi/2)%pi - pi/2          # (n)
    dtheta2 = regular_theta(gtheta + pi / 2)  # (gtheta+pi)% pi -pi/2    # (n)
    abs_dtheta1 = torch.abs(dtheta1)
    abs_dtheta2 = torch.abs(dtheta2)
    dtheta = torch.where(abs_dtheta1 < abs_dtheta2, dtheta1,
                         dtheta2)  # dtheta (n)     # (-pi/4,pi/4):gtheta (pi/4,pi/2):pi/2-gtheta (-pi/2,-pi/4): pi/2+gtheta  # 这里才最终得到所谓的
    bboxes = torch.stack([gx, gy, gw, gh, gtheta.unsqueeze(-1)], dim=-1)
    bboxes = regular_obb(bboxes)
    return bboxes.squeeze(1)
